import jieba
from sklearn.feature_extraction.text import TfidfVectorizer

import pandas as pd

from joblib import load


def load_data(file_path):
    return pd.read_csv(file_path)


def generate_ones_array(label: int, size: int):
    return [label] * size


def read_txt_file_to_list(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
    lines = [line.strip() for line in lines]
    return lines


def init_data() -> (list, list):
    global texts
    data = load_data('Detect_data/train-data-all-0907-001-utf8.csv')
    data['content'] = data['content'].fillna('')
    texts = data['content'].tolist()
    data['content'] = data['content'].fillna(0)
    labels = data['label'].tolist()
    lines_list = read_txt_file_to_list('Detect_data/attack-text.txt')
    texts = texts + lines_list
    lines_labels = generate_ones_array(1, len(lines_list))
    labels = labels + lines_labels
    lines_list = read_txt_file_to_list('Detect_data/normal-data-text.txt')
    texts = texts + lines_list
    lines_labels = generate_ones_array(0, len(lines_list))
    labels = labels + lines_labels
    return texts, labels


if __name__ == '__main__':
    texts, labels = init_data()
    vectorizer = TfidfVectorizer(tokenizer=jieba.lcut)
    vectorizer.fit_transform(texts)
    loaded_classifier = load('fit-model.joblib')

    # X_test = vectorizer.transform(["测试一下"])
    # # 使用加载的模型进行预测
    # y_pred = loaded_classifier.predict(X_test)
    # print(y_pred)

    data = load_data('Detect_data/train-data-all-0907-001-utf8.csv')
    data['content'] = data['content'].fillna('')
    texts = data['content'].tolist()
    for test_sentence in texts:
        # TODO 解析测试的content，根据现有模型，获取预测结果label
        X_test = vectorizer.transform([test_sentence])
        # 使用加载的模型进行预测
        y_pred = loaded_classifier.predict(X_test)
        print(y_pred)
